# 导入 MXNet 模块
import mxnet as mx
from mxnet import gluon, autograd, nd
from mxnet.gluon import nn

# 设置随机种子（确保结果可复现）
mx.random.seed(42)

# 定义简单的神经网络模型
class SimpleNet(nn.Block):
    def __init__(self, **kwargs):
        super(SimpleNet, self).__init__(**kwargs)
        self.fc1 = nn.Dense(128)  # 全连接层，128 个神经元
        self.fc2 = nn.Dense(64)   # 全连接层，64 个神经元
        self.fc3 = nn.Dense(10)   # 输出层，10 个神经元（假设是 10 分类任务）

    def forward(self, x):
        x = nd.relu(self.fc1(x))  # 激活函数 ReLU
        x = nd.relu(self.fc2(x))
        return self.fc3(x)

# 创建模型实例
net = SimpleNet()

# 初始化模型参数
net.initialize(mx.init.Xavier())

# 定义损失函数（交叉熵损失）
loss_fn = gluon.loss.SoftmaxCrossEntropyLoss()

# 定义优化器（随机梯度下降）
trainer = gluon.Trainer(net.collect_params(), 'sgd', {'learning_rate': 0.01})

# 生成示例数据
batch_size = 32
X = nd.random.normal(shape=(batch_size, 784))  # 假设输入是 784 维向量（如 MNIST 图像）
Y = nd.random.randint(0, 10, shape=(batch_size,))  # 假设标签是 0 到 9 的整数

# 训练模型
epochs = 5
for epoch in range(epochs):
    with autograd.record():
        output = net(X)  # 前向传播
        loss = loss_fn(output, Y)  # 计算损失
    loss.backward()  # 反向传播
    trainer.step(batch_size)  # 更新参数
    print(f"Epoch [{epoch+1}/{epochs}], Loss: {loss.mean().asnumpy()}")

# 测试模型
test_X = nd.random.normal(shape=(batch_size, 784))
preds = net(test_X).argmax(axis=1)  # 获取预测类别
print("预测结果：", preds.asnumpy())